import os
import argparse
import torch
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import sys

# 添加项目根目录到路径
sys.path.append(os.path.abspath(os.path.dirname(os.path.dirname(__file__))))

from models.colorization_model import ColorizationGenerator
from utils.data_utils import prepare_grayscale_image, lab_to_rgb

def colorize_image(model_path, image_path, output_path=None, img_size=256, use_cuda=True):
    """
    对单张灰度图像进行着色
    
    参数:
        model_path (str): 模型路径
        image_path (str): 输入图像路径
        output_path (str, optional): 输出图像路径
        img_size (int): 图像大小
        use_cuda (bool): 是否使用CUDA
    
    返回:
        output_image (PIL.Image): 着色后的图像
    """
    # 设置设备
    device = torch.device('cuda' if torch.cuda.is_available() and use_cuda else 'cpu')
    print(f"使用设备: {device}")
    
    # 加载模型
    model = ColorizationGenerator().to(device)
    checkpoint = torch.load(model_path, map_location=device)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    print(f"模型已加载: {model_path}")
    
    # 准备输入图像
    L = prepare_grayscale_image(image_path, img_size=img_size).to(device)
    print(f"图像已加载: {image_path}")
    
    # 推理
    with torch.no_grad():
        ab_pred = model(L)
    print("图像着色完成")
    
    # 转换为RGB
    L_np = L.cpu()
    ab_np = ab_pred.cpu()
    rgb_image = lab_to_rgb(L_np, ab_np)[0]
    
    # 转换为PIL图像
    output_image = Image.fromarray((rgb_image * 255).astype(np.uint8))
    
    # 保存结果
    if output_path:
        output_image.save(output_path)
        print(f"结果已保存: {output_path}")
    
    return output_image

def visualize_result(original_path, colorized_path, output_path=None):
    """
    可视化原始图像和着色结果
    
    参数:
        original_path (str): 原始图像路径
        colorized_path (str): 着色后图像路径
        output_path (str, optional): 输出图像路径
    """
    # 读取图像
    original = Image.open(original_path).convert('L')  # 转为灰度图
    colorized = Image.open(colorized_path)
    
    # 创建图像
    fig, axes = plt.subplots(1, 2, figsize=(12, 6))
    
    # 显示原始灰度图
    axes[0].imshow(original, cmap='gray')
    axes[0].set_title('原始灰度图')
    axes[0].axis('off')
    
    # 显示着色后图像
    axes[1].imshow(colorized)
    axes[1].set_title('着色后图像')
    axes[1].axis('off')
    
    plt.tight_layout()
    
    # 保存结果
    if output_path:
        plt.savefig(output_path)
        print(f"对比图已保存: {output_path}")
    
    plt.show()

def main():
    parser = argparse.ArgumentParser(description='灰度图像着色推理')
    
    parser.add_argument('--model', type=str, required=True, help='模型路径')
    parser.add_argument('--input', type=str, required=True, help='输入图像路径')
    parser.add_argument('--output', type=str, default=None, help='输出图像路径')
    parser.add_argument('--img_size', type=int, default=256, help='图像大小')
    parser.add_argument('--no_cuda', action='store_true', help='不使用CUDA')
    parser.add_argument('--visualize', action='store_true', help='可视化结果')
    
    args = parser.parse_args()
    
    # 着色图像
    output_path = args.output
    if output_path is None:
        # 如果未指定输出路径，则在输入图像的同一目录下创建
        input_dir = os.path.dirname(args.input)
        input_name = os.path.basename(args.input)
        output_name = f"colorized_{input_name}"
        output_path = os.path.join(input_dir, output_name)
    
    colorize_image(
        args.model,
        args.input,
        output_path,
        img_size=args.img_size,
        use_cuda=not args.no_cuda
    )
    
    # 可视化结果
    if args.visualize:
        vis_output_path = os.path.splitext(output_path)[0] + "_comparison.png"
        visualize_result(args.input, output_path, vis_output_path)

if __name__ == '__main__':
    main() 